from typing import Dict, Optional
from fastapi import FastAPI, HTTPException, Depends, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from pydantic import BaseModel
import secrets
from tqsdk import TqApi, TqAuth, TqKq
import asyncio
import uvicorn
from fastapi.staticfiles import StaticFiles
from fastapi.responses import RedirectResponse
import os
from fastapi.middleware.cors import CORSMiddleware

# 创建静态文件目录（如果不存在）
os.makedirs("static/css", exist_ok=True)
os.makedirs("static/js", exist_ok=True)

app = FastAPI(title="TQSDK账户查询API", description="提供TQSDK账户资金和持仓查询功能")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 安全认证
security = HTTPBasic()

# 用户凭据模型
class UserCredentials(BaseModel):
    username: str
    password: str

# 账户信息响应模型
class AccountInfo(BaseModel):
    available: float  # 可用资金
    balance: float    # 账户权益
    currency: str     # 币种
    margin: float     # 占用保证金
    risk_ratio: float # 风险度
    pre_balance: float # 昨日账户权益
    float_profit: float # 浮动盈亏
    close_profit: float # 平仓盈亏

# 持仓信息模型
class PositionInfo(BaseModel):
    symbol: str       # 合约代码
    pos_long: int     # 多头持仓手数
    pos_short: int    # 空头持仓手数
    volume_long_today: int  # 今日多头开仓手数
    volume_short_today: int # 今日空头开仓手数
    float_profit_long: float # 多头浮动盈亏
    float_profit_short: float # 空头浮动盈亏
    position_price_long: float # 多头持仓均价
    position_price_short: float # 空头持仓均价

# 响应模型
class ApiResponse(BaseModel):
    success: bool
    message: str
    data: Optional[Dict] = None

# 验证用户凭据
def verify_credentials(credentials: HTTPBasicCredentials = Depends(security)):
    # 这里应该实现真实的用户验证逻辑
    # 示例中使用简单的硬编码验证
    correct_username = "admin"
    correct_password = "password"
    
    is_username_correct = secrets.compare_digest(credentials.username, correct_username)
    is_password_correct = secrets.compare_digest(credentials.password, correct_password)
    
    if not (is_username_correct and is_password_correct):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid username or password",
            headers={"WWW-Authenticate": "Basic"},
        )
    
    return credentials.username

# 查询TQSDK账户信息
async def query_tqsdk_account(username: str, password: str):
    try:
        # 打印调试信息
        print(f"正在使用用户名: {username} 和密码: {'*' * len(password)} 连接TQSDK")
        
        # 创建一个事件循环
        loop = asyncio.get_event_loop()
        
        # 在线程池中执行阻塞操作
        def run_tqsdk_query():
            api = None
            try:
                # 使用提供的用户名和密码连接TQSDK
                api = TqApi(TqKq(), auth=TqAuth(username, password))
                
                # 获取账户信息
                account_info = api.get_account()
                
                # 获取持仓信息
                # 等待数据更新，设置10秒超时
                if not api.wait_update():
                    raise TimeoutError("获取账户数据超时，请检查网络连接或TQSDK服务")
                
                # 构建持仓数据
                position_data = []
                # 获取所有持仓合约
                positions = api.get_position()
                for position in positions.values():
                    if position.pos_long > 0 or position.pos_short > 0:  # 只添加有持仓的合约
                        position_data.append({
                            "symbol": position.instrument_id,
                            "pos_long": position.pos_long,
                            "pos_short": position.pos_short,
                            "volume_long_today": position.volume_long_today,
                            "volume_short_today": position.volume_short_today,
                            "float_profit_long": position.float_profit_long,
                            "float_profit_short": position.float_profit_short,
                            "position_price_long": position.position_price_long,
                            "position_price_short": position.position_price_short
                        })
                
                # 返回结果
                return {
                    "account": {
                        "available": account_info.available,
                        "balance": account_info.balance,
                        "currency": account_info.currency,
                        "margin": account_info.margin,
                        "risk_ratio": account_info.risk_ratio,
                        "pre_balance": account_info.pre_balance,
                        "float_profit": account_info.float_profit,
                        "close_profit": account_info.close_profit
                    },
                    "positions": position_data
                }
            except Exception as e:
                return {"error": str(e)}
            finally:
                # 确保API连接关闭
                if api is not None:
                    api.close()
        
        # 在线程池中执行TQSDK查询
        result = await loop.run_in_executor(None, run_tqsdk_query)
        
        if "error" in result:
            return {"success": False, "message": result["error"], "data": None}
        
        return {"success": True, "message": "查询成功", "data": result}
    except (Exception, asyncio.CancelledError) as e:
        return {"success": False, "message": f"查询被取消或失败: {str(e)}", "data": None}

@app.post("/get_account_info", response_model=ApiResponse)
async def get_account_info(user: UserCredentials):
    """
    获取账户资金和持仓信息
    
    - **username**: TQSDK用户名
    - **password**: TQSDK密码
    
    返回账户资金和持仓信息
    """
    result = await query_tqsdk_account(user.username, user.password)
    return result

# 挂载静态文件目录
app.mount("/static", StaticFiles(directory="static"), name="static")

@app.get("/")
async def root():
    """重定向到Web界面"""
    return RedirectResponse(url="/static/index.html")

if __name__ == "__main__":
    try:
        config = uvicorn.Config(app, host="127.0.0.1", port=8000, timeout_keep_alive=30)
        server = uvicorn.Server(config)
        server.run()
    except KeyboardInterrupt:
        print("\n服务器已停止")
    finally:
        # 显式关闭事件循环（如果存在）
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                loop.stop()
                loop.close()
        except RuntimeError:
            pass